LLM Based Task#2283
Conversation
Signed-off-by: Kyle Zheng <126034466+KyleZheng1284@users.noreply.github.com>
Greptile SummaryThis PR introduces a reusable
|
| Filename | Overview |
|---|---|
| nemo_retriever/src/nemo_retriever/models/llm/tasks/rag_answer.py | Encapsulates RAG prompt construction and reasoning controls into a typed task; custom prompt template is not validated at construction time, leading to silent request_error at inference for malformed templates. |
| nemo_retriever/src/nemo_retriever/models/llm/tasks/base.py | New GenerationTask ABC with strict invoke() lifecycle and compatibility execute() facade; bare excepts in invoke() at _preflight_error, build_request, and parse phases swallow root exceptions without logging (previously flagged). |
| nemo_retriever/src/nemo_retriever/operators/generation/base.py | New TextGenerationOperator with concurrent row execution, position-tracking, and safe graph serialization; bare except in _execute_row and _failure_model lack exc_info=True (previously flagged); process() logic is otherwise solid. |
| nemo_retriever/src/nemo_retriever/models/llm/tasks/generic.py | New GenericPromptTask with thorough template validation; imports private _apply_reasoning_control from sibling rag_answer module, and uses string comparison for exception dispatch which is brittle. |
| nemo_retriever/src/nemo_retriever/common/params/models.py | Adds TextGenerationParams, LLMSamplingOverrides, and environment-reference credential handling; secret redaction helpers are thorough; LLMRemoteClientParams now defers api_key resolution to worker side safely. |
| nemo_retriever/src/nemo_retriever/tools/evaluation/generation.py | QAGenerationOperator refactored to extend TextGenerationOperator; backward-compatible flat constructor preserved; legacy LLMClient adapter correctly forwards reasoning_enabled; overwrite=True hardcoded for backward compatibility. |
| nemo_retriever/src/nemo_retriever/models/llm/clients/litellm.py | Refactored to delegate to RagAnswerTask; new _call_complete validates response shape with UnsupportedTextResponseError; environment-reference api_key resolved at call time correctly. |
| nemo_retriever/src/nemo_retriever/graph/graph_pipeline_registry.py | Significant expansion of graph serialization to version 2 with comprehensive secret field detection and environment-reference encoding; _encode_value covers Pydantic models, dicts, lists, and tuples safely. |
| nemo_retriever/src/nemo_retriever/models/llm/types.py | Adds TextCompletionClient protocol, UnsupportedTextResponseError, GenerationRequest (with snapshot-on-construct immutability), and GeneratedTextResult; clean design with appropriate validation in post_init. |
Sequence Diagram
%%{init: {'theme': 'neutral'}}%%
sequenceDiagram
participant Op as TextGenerationOperator
participant Pool as ThreadPoolExecutor
participant Row as _execute_row
participant Task as GenerationTask.invoke()
participant Client as TextCompletionClient
Op->>Pool: submit _execute_row(position, inputs) per row
Pool->>Row: execute
Row->>Task: _execute_task(inputs)
Task->>Task: "_preflight_error(**inputs)"
Task->>Task: "build_request(**inputs) → GenerationRequest"
Task->>Client: complete(messages, max_tokens, extra_params)
Client-->>Task: (raw_text, latency_s)
Task->>Task: parse(raw_text) → text
Task-->>Row: GeneratedTextResult
Row-->>Pool: (position, result)
Pool-->>Op: as_completed futures
Op->>Op: "assert result_position == position"
Op->>Op: write output columns to DataFrame copy
%%{init: {'theme': 'base', 'themeVariables': {"darkMode": true, "background": "#0d1117", "primaryColor": "#21262d", "primaryTextColor": "#e6edf3", "primaryBorderColor": "#8b949e", "lineColor": "#8b949e", "textColor": "#e6edf3", "edgeLabelBackground": "#161b22", "actorBkg": "#21262d", "actorBorder": "#8b949e", "actorTextColor": "#e6edf3", "actorLineColor": "#8b949e", "signalColor": "#8b949e", "signalTextColor": "#e6edf3", "noteBkgColor": "#373320", "noteBorderColor": "#d4a72c", "noteTextColor": "#f0e6c0", "labelBoxBkgColor": "#21262d", "labelBoxBorderColor": "#8b949e", "labelTextColor": "#e6edf3", "loopTextColor": "#e6edf3", "activationBkgColor": "#30363d", "activationBorderColor": "#8b949e"}}}%%
sequenceDiagram
participant Op as TextGenerationOperator
participant Pool as ThreadPoolExecutor
participant Row as _execute_row
participant Task as GenerationTask.invoke()
participant Client as TextCompletionClient
Op->>Pool: submit _execute_row(position, inputs) per row
Pool->>Row: execute
Row->>Task: _execute_task(inputs)
Task->>Task: "_preflight_error(**inputs)"
Task->>Task: "build_request(**inputs) → GenerationRequest"
Task->>Client: complete(messages, max_tokens, extra_params)
Client-->>Task: (raw_text, latency_s)
Task->>Task: parse(raw_text) → text
Task-->>Row: GeneratedTextResult
Row-->>Pool: (position, result)
Pool-->>Op: as_completed futures
Op->>Op: "assert result_position == position"
Op->>Op: write output columns to DataFrame copy
Prompt To Fix All With AI
Fix the following 1 code review issue. Work through them one at a time, proposing concise fixes.
---
### Issue 1 of 1
nemo_retriever/src/nemo_retriever/models/llm/tasks/rag_answer.py:135-137
**Custom `prompt` template not validated at construction time**
`self.prompt.format(context=context, query=query)` runs at inference time without any prior validation of the template. If the template contains undeclared placeholders — e.g. `RagAnswerTask(prompt="{typo}")` — Python raises `KeyError: 'typo'` inside `build_request`. That exception is silently caught by `invoke`'s `build_request` handler and converted to `request_error` with no diagnostic message about the template, leaving every row in a batch failing with an opaque error.
`SummarizeTask` demonstrates the correct pattern: it calls `_summary_prompt_fields(self.prompt)` in `__post_init__` so an invalid template surfaces immediately at task construction. `RagAnswerTask` should do the same — validate at `__post_init__` that `self.prompt` uses only the `{context}` and `{query}` placeholders.
Reviews (3): Last reviewed commit: "Address generation review feedback" | Re-trigger Greptile
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| # SPDX-FileCopyrightText: Copyright (c) 2024-26, NVIDIA CORPORATION & AFFILIATES. | ||
| # All rights reserved. | ||
| # SPDX-License-Identifier: Apache-2.0 |
There was a problem hiding this comment.
The tasks/ package and operator files use three different year formats: 2024-26 (non-standard shorthand), 2024-25, and 2026 (standalone). The spdx-license-header rule requires the current year. New files in this PR use 2024-26 (tasks/__init__.py, tasks/base.py, tasks/generic.py, tasks/rag_answer.py, tasks/summarize.py) while test_generation_tasks.py uses 2026. The standard form in the rest of the repo is the four-digit year; the abbreviated 2024-26 is not conventional. All new files should use the same format as the existing codebase.
Prompt To Fix With AI
This is a comment left during a code review.
Path: nemo_retriever/src/nemo_retriever/models/llm/tasks/__init__.py
Line: 1-3
Comment:
**SPDX copyright year format inconsistency across new files**
The `tasks/` package and operator files use three different year formats: `2024-26` (non-standard shorthand), `2024-25`, and `2026` (standalone). The `spdx-license-header` rule requires the current year. New files in this PR use `2024-26` (`tasks/__init__.py`, `tasks/base.py`, `tasks/generic.py`, `tasks/rag_answer.py`, `tasks/summarize.py`) while `test_generation_tasks.py` uses `2026`. The standard form in the rest of the repo is the four-digit year; the abbreviated `2024-26` is not conventional. All new files should use the same format as the existing codebase.
How can I resolve this? If you propose a fix, please make it concise.Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!
Signed-off-by: Kyle Zheng <126034466+KyleZheng1284@users.noreply.github.com>
Signed-off-by: Kyle Zheng <126034466+KyleZheng1284@users.noreply.github.com>
| if self.prompt is not None: | ||
| context = "\n\n---\n\n".join(chunks) if chunks else "(no context retrieved)" | ||
| messages[-1]["content"] = self.prompt.format(context=context, query=query) |
There was a problem hiding this comment.
prompt template not validated at construction time
self.prompt.format(context=context, query=query) runs at inference time without any prior validation of the template. If the template contains undeclared placeholders — e.g. RagAnswerTask(prompt="{typo}") — Python raises KeyError: 'typo' inside build_request. That exception is silently caught by invoke's build_request handler and converted to request_error with no diagnostic message about the template, leaving every row in a batch failing with an opaque error.
SummarizeTask demonstrates the correct pattern: it calls _summary_prompt_fields(self.prompt) in __post_init__ so an invalid template surfaces immediately at task construction. RagAnswerTask should do the same — validate at __post_init__ that self.prompt uses only the {context} and {query} placeholders.
Prompt To Fix With AI
This is a comment left during a code review.
Path: nemo_retriever/src/nemo_retriever/models/llm/tasks/rag_answer.py
Line: 135-137
Comment:
**Custom `prompt` template not validated at construction time**
`self.prompt.format(context=context, query=query)` runs at inference time without any prior validation of the template. If the template contains undeclared placeholders — e.g. `RagAnswerTask(prompt="{typo}")` — Python raises `KeyError: 'typo'` inside `build_request`. That exception is silently caught by `invoke`'s `build_request` handler and converted to `request_error` with no diagnostic message about the template, leaving every row in a batch failing with an opaque error.
`SummarizeTask` demonstrates the correct pattern: it calls `_summary_prompt_fields(self.prompt)` in `__post_init__` so an invalid template surfaces immediately at task construction. `RagAnswerTask` should do the same — validate at `__post_init__` that `self.prompt` uses only the `{context}` and `{query}` placeholders.
How can I resolve this? If you propose a fix, please make it concise.1 participant
Description
Checklist